import tqdm
import numpy as np
from Config import *
import torch
from model import *
if __name__ == '__main__':

    print('prediction')

    content = open('/Users/Downloads/burst_raw/competition_test_input.0.2.bin', 'rb').read()
    samples_ref = np.frombuffer(content, dtype = 'uint16').reshape((-1,256,256))
    fout = open('/Users/Downloads/burst_raw/competition_prediction.0.2.bin', 'wb')
    model = Predictor()
    for i in tqdm.tqdm(range(0, len(samples_ref), Config.batch_size)):
        i_end = min(i + Config.batch_size, len(samples_ref))
        batch_inp = torch.tensor(np.float32(samples_ref[i:i_end, None, :, :]) * np.float32(1 / 65536))
        pred = model(batch_inp)

        pred = (pred.numpy()[:, 0, :, :] * 65536).clip(0, 65535).astype('uint16')
        fout.write(pred.tobytes())

    fout.close()